import os
import time
import json
import uuid
from typing import Optional

import requests
from flask import Flask, request, jsonify, make_response, g, send_from_directory
from PIL import Image
from io import BytesIO
import logging
import traceback
from werkzeug.utils import secure_filename


DEFAULT_BASE_URL = "https://api-inference.modelscope.cn/"

# 配置上传
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'webp'}
MAX_CONTENT_LENGTH = 16 * 1024 * 1024  # 16MB

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH

# 确保上传目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# Configure logging to always emit to console and honor LOG_LEVEL env
def _configure_logging():
    log_level_name = os.getenv("LOG_LEVEL", "INFO").upper()
    log_level = getattr(logging, log_level_name, logging.INFO)
    root_logger = logging.getLogger()
    root_logger.setLevel(log_level)

    # Ensure there is a console handler with a clear format
    if not any(isinstance(h, logging.StreamHandler) for h in root_logger.handlers):
        console_handler = logging.StreamHandler()
        console_handler.setLevel(log_level)
        console_handler.setFormatter(logging.Formatter(
            fmt="%(asctime)s %(levelname)s %(name)s - %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        ))
        root_logger.addHandler(console_handler)

    # Align Flask/werkzeug loggers
    app.logger.setLevel(log_level)
    werkzeug_logger = logging.getLogger("werkzeug")
    werkzeug_logger.setLevel(log_level)

_configure_logging()


def _json_error(status_code: int, detail: str):
    return make_response(jsonify({"detail": detail}), status_code)


def _get_api_key() -> str:
    api_key = os.getenv("MODELSCOPE_API_KEY","ms-e0c47d3f-1a5c-415a-9daa-d8ec518e8266") or os.getenv("API_KEY","ms-e0c47d3f-1a5c-415a-9daa-d8ec518e8266")
    if not api_key:
        raise RuntimeError("缺少环境变量 MODELSCOPE_API_KEY 或 API_KEY")
    return api_key


def _get_base_url() -> str:
    return (os.getenv("MODELSCOPE_BASE_URL") or DEFAULT_BASE_URL).rstrip("/") + "/"


def _get_auth_salt() -> str:
    auth_salt = os.getenv("AUTH_SALT","GeovisEarth")
    if not auth_salt:
        raise RuntimeError("缺少环境变量 AUTH_SALT 用于盐认证")
    return auth_salt


def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def get_file_url(filename):
    """获取文件的HTTP访问URL"""
    # 获取请求的主机和端口
    host = request.host
    scheme = request.scheme
    return f"{scheme}://{host}/uploads/{filename}"


def compute_size_from_image(image_url: str) -> str:
    resp = requests.get(image_url, timeout=30)
    resp.raise_for_status()
    with Image.open(BytesIO(resp.content)) as im:
        w, h = im.size
    if w > 2048 or h > 2048:
        return "2048x2048"
    return f"{w}x{h}"


def submit_generation(api_key: str, base_url: str, payload: dict) -> str:
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "X-ModelScope-Async-Mode": "true",
    }
    resp = requests.post(
        f"{base_url}v1/images/generations",
        headers=headers,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        timeout=60,
    )
    resp.raise_for_status()
    data = resp.json()
    if "task_id" not in data:
        raise RuntimeError(f"提交任务失败: {data}")
    return data["task_id"]


def poll_task(api_key: str, base_url: str, task_id: str, poll_interval: float, timeout_seconds: float) -> str:
    headers = {
        "Authorization": f"Bearer {api_key}",
        "X-ModelScope-Task-Type": "image_generation",
    }
    deadline = time.time() + timeout_seconds
    while True:
        if time.time() > deadline:
            raise TimeoutError("生成超时")
        r = requests.get(f"{base_url}v1/tasks/{task_id}", headers=headers, timeout=60)
        r.raise_for_status()
        data = r.json()
        status = data.get("task_status")
        if status == "SUCCEED":
            images = data.get("output_images") or []
            if not images:
                raise RuntimeError("任务成功但未返回图片")
            return images[0]
        if status == "FAILED":
            raise RuntimeError(f"生成失败: {data}")
        time.sleep(poll_interval)


def submit_image_generation(api_key: str, base_url: str, payload: dict) -> str:
    """提交图像生成任务"""
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "X-ModelScope-Async-Mode": "true",
    }
    resp = requests.post(
        f"{base_url}v1/images/generations",
        headers=headers,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        timeout=60,
    )
    resp.raise_for_status()
    data = resp.json()
    if "task_id" not in data:
        raise RuntimeError(f"提交任务失败: {data}")
    return data["task_id"]


def poll_image_generation_task(api_key: str, base_url: str, task_id: str, poll_interval: float, timeout_seconds: float) -> str:
    """轮询图像生成任务状态"""
    headers = {
        "Authorization": f"Bearer {api_key}",
        "X-ModelScope-Task-Type": "image_generation",
    }
    deadline = time.time() + timeout_seconds
    while True:
        if time.time() > deadline:
            raise TimeoutError("生成超时")
        r = requests.get(f"{base_url}v1/tasks/{task_id}", headers=headers, timeout=60)
        r.raise_for_status()
        data = r.json()
        status = data.get("task_status")
        if status == "SUCCEED":
            images = data.get("output_images") or []
            if not images:
                raise RuntimeError("任务成功但未返回图片")
            return images[0]
        if status == "FAILED":
            raise RuntimeError(f"生成失败: {data}")
        time.sleep(poll_interval)


@app.route("/healthz", methods=["GET"])
def healthz():
    return jsonify({"status": "ok"})


@app.route("/upload", methods=["POST"])
def upload_file():
    """文件上传端点"""
    try:
        # 简单盐认证
        expected_salt = _get_auth_salt()
        provided_salt = request.headers.get("X-Auth-Salt")
        if not provided_salt or provided_salt != expected_salt:
            return _json_error(401, "未授权: 盐认证失败")
        
        # 检查是否有文件
        if 'file' not in request.files:
            return _json_error(400, "没有文件被上传")
        
        file = request.files['file']
        if file.filename == '':
            return _json_error(400, "没有选择文件")
        
        if file and allowed_file(file.filename):
            # 生成安全的文件名
            original_filename = secure_filename(file.filename)
            file_extension = original_filename.rsplit('.', 1)[1].lower() if '.' in original_filename else ''
            unique_filename = f"{uuid.uuid4().hex}.{file_extension}"
            
            # 保存文件
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
            file.save(file_path)
            
            # 获取文件URL
            file_url = get_file_url(unique_filename)
            
            # 获取图片分辨率
            try:
                with Image.open(file_path) as img:
                    width, height = img.size
                    image_dimensions = f"{width}x{height}"
            except Exception as e:
                app.logger.warning(f"无法获取图片分辨率: {e}")
                image_dimensions = "未知"
            
            app.logger.info(f"文件上传成功: {unique_filename}, 分辨率: {image_dimensions}")
            
            return jsonify({
                "success": True,
                "filename": unique_filename,
                "original_filename": original_filename,
                "file_url": file_url,
                "size": os.path.getsize(file_path),
                "dimensions": image_dimensions,
                "width": width,
                "height": height
            })
        else:
            return _json_error(400, f"不支持的文件类型。允许的类型: {', '.join(ALLOWED_EXTENSIONS)}")
            
    except Exception as e:
        app.logger.exception("文件上传失败")
        return _json_error(500, f"文件上传失败: {str(e)}")


@app.route("/uploads/<filename>")
def uploaded_file(filename):
    """提供上传文件的访问"""
    try:
        return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
    except FileNotFoundError:
        return _json_error(404, "文件不存在")


# Basic request/response logging
@app.before_request
def _log_request_start():
    g._req_start_ts = time.time()
    try:
        body = request.get_json(silent=True)
    except Exception:
        body = None
    app.logger.info(f"--> {request.method} {request.path} headers={{'X-Auth-Salt': '{request.headers.get('X-Auth-Salt', '')}'}} body={body}")


@app.route("/generate", methods=["POST"])
def generate():
    try:
        api_key = _get_api_key()
    except RuntimeError as e:
        app.logger.error(f"配置错误: {e}")
        return _json_error(500, str(e))

    base_url = _get_base_url()
    # 简单盐认证（固定值比对）
    try:
        expected_salt = _get_auth_salt()
    except RuntimeError as e:
        app.logger.error(f"配置错误: {e}")
        return _json_error(500, str(e))
    provided_salt = request.headers.get("X-Auth-Salt")
    if not provided_salt or provided_salt != expected_salt:
        return _json_error(401, "未授权: 盐认证失败")

    data = request.get_json(silent=True)
    if not data:
        return _json_error(400, "请求体不是有效的 JSON")

    model = "Qwen/Qwen-Image-Edit"
    prompt = data.get("prompt")
    image_url = data.get("image_url")
    negative_prompt = data.get("negative_prompt") or "lowres, watermark, text artifacts"
    size = data.get("size")
    seed = data.get("seed", 12345)
    steps = data.get("steps", 50)
    guidance = data.get("guidance", 3.5)
    poll_interval_seconds = float(data.get("poll_interval_seconds", 3.0))
    timeout_seconds = float(data.get("timeout_seconds", 180.0))

    if not prompt or not image_url:
        return _json_error(400, "缺少必填字段: prompt, image_url")

    target_size = size
    if not target_size:
        try:
            target_size = compute_size_from_image(str(image_url))
        except Exception as e:
            app.logger.exception("计算尺寸失败")
            return _json_error(400, f"计算尺寸失败: {e}")

    payload = {
        "model": model,
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "size": target_size,
        "seed": seed,
        "steps": steps,
        "guidance": guidance,
        "image_url": str(image_url),
    }

    try:
        task_id = submit_generation(api_key, base_url, payload)
        output_image_url = poll_task(
            api_key,
            base_url,
            task_id,
            poll_interval=poll_interval_seconds,
            timeout_seconds=timeout_seconds,
        )
    except requests.HTTPError as e:
        app.logger.exception("上游 HTTP 错误")
        return _json_error(502, f"上游错误: {e}")
    except TimeoutError as e:
        app.logger.warning(f"超时: {e}")
        return _json_error(504, str(e))
    except Exception as e:
        app.logger.exception("处理失败")
        return _json_error(502, f"处理失败: {e}")

    return jsonify({
        "task_id": task_id,
        "output_image_url": output_image_url,
        "size": target_size,
    })


@app.route("/generate_image", methods=["POST"])
def generate_image():
    """图像生成端点 - 从文本生成图像"""
    try:
        api_key = _get_api_key()
    except RuntimeError as e:
        app.logger.error(f"配置错误: {e}")
        return _json_error(500, str(e))

    base_url = _get_base_url()
    # 简单盐认证（固定值比对）
    try:
        expected_salt = _get_auth_salt()
    except RuntimeError as e:
        app.logger.error(f"配置错误: {e}")
        return _json_error(500, str(e))
    provided_salt = request.headers.get("X-Auth-Salt")
    if not provided_salt or provided_salt != expected_salt:
        return _json_error(401, "未授权: 盐认证失败")

    data = request.get_json(silent=True)
    if not data:
        return _json_error(400, "请求体不是有效的 JSON")

    model = "Qwen/Qwen-Image"
    prompt = data.get("prompt")
    negative_prompt = data.get("negative_prompt") or "lowres, watermark, text artifacts"
    size = data.get("size", "1024x1024")
    seed = data.get("seed", 12345)
    steps = data.get("steps", 50)
    guidance = data.get("guidance", 3.5)
    poll_interval_seconds = float(data.get("poll_interval_seconds", 5.0))
    timeout_seconds = float(data.get("timeout_seconds", 300.0))

    if not prompt:
        return _json_error(400, "缺少必填字段: prompt")

    payload = {
        "model": model,
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "size": size,
        "seed": seed,
        "steps": steps,
        "guidance": guidance,
    }

    try:
        task_id = submit_image_generation(api_key, base_url, payload)
        output_image_url = poll_image_generation_task(
            api_key,
            base_url,
            task_id,
            poll_interval=poll_interval_seconds,
            timeout_seconds=timeout_seconds,
        )
    except requests.HTTPError as e:
        app.logger.exception("上游 HTTP 错误")
        return _json_error(502, f"上游错误: {e}")
    except TimeoutError as e:
        app.logger.warning(f"超时: {e}")
        return _json_error(504, str(e))
    except Exception as e:
        app.logger.exception("处理失败")
        return _json_error(502, f"处理失败: {e}")

    return jsonify({
        "task_id": task_id,
        "output_image_url": output_image_url,
        "size": size,
        "model": model,
    })


@app.errorhandler(Exception)
def handle_unhandled_exception(e: Exception):
    app.logger.exception("未捕获的服务器异常")
    return _json_error(500, "服务器内部错误" )


# CORS支持和请求日志记录
@app.after_request
def after_request(response):
    # 添加CORS头
    response.headers.add('Access-Control-Allow-Origin', '*')
    response.headers.add('Access-Control-Allow-Headers', 'Content-Type, X-Auth-Salt')
    response.headers.add('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
    
    # 记录请求日志
    try:
        duration_ms = int((time.time() - getattr(g, "_req_start_ts", time.time())) * 1000)
    except Exception:
        duration_ms = -1
    app.logger.info(f"<-- {request.method} {request.path} {response.status_code} {duration_ms}ms")
    
    return response

# 处理OPTIONS预检请求
@app.route('/generate', methods=['OPTIONS'])
@app.route('/generate_image', methods=['OPTIONS'])
@app.route('/upload', methods=['OPTIONS'])
def handle_options():
    return '', 200

if __name__ == "__main__":
    # 本地开发启动：python service.py
    host = os.getenv("HOST", "0.0.0.0")
    debug = os.getenv("FLASK_DEBUG", "0") == "1"
    port = int(os.getenv("PORT", "5001"))
    app.run(host=host, port=port, debug=debug)
    

